from invfun import inv, inv_1
from onepair import BasicGears, Gear, PlanetaryGear, HiddenPrints, PrintToTxt
import math
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits import mplot3d

if __name__ == '__main__':
    # 齿轮和插齿刀基本参数
    basegear = BasicGears(20, 1, 0.25, 58, 0, 1.25)
    hh1 = 0.001  # 遍历的间隔
    xiubuflag = False
    # 齿数基本关系 为了得到理想的啮合角，行星轮齿数比标准齿轮减1
    m = 1.75
    z1 = 18
    z2 = 108
    z3 = 44
    z_segma_13 = z1+z3
    z_segma_32 = z2-z3
    alpha = math.radians(20)

    # pic 1
    # 先给定一个啮合角，以x2为变量作目标函数二维图
    alpha_work_32 = math.radians(20.66)

    # 根据齿数和关系计算另一个啮合角
    alpha_work_13 = math.acos(z_segma_13/z_segma_32*math.cos(alpha_work_32))

    # 计算两组变位系数和
    x_segma_32 = z_segma_32 * \
        (inv(alpha_work_32)-inv(alpha)) / (2*math.tan(alpha))

    x_segma_13 = z_segma_13 * \
        (inv(alpha_work_13)-inv(alpha)) / (2*math.tan(alpha))

    x2_range = np.arange(0, 1, hh1)

    x1_list = []
    x3_list = []
    x2_list = []

    yeta13_1_list = []
    yeta13_3_list = []
    yeta32_2_list = []
    yeta32_3_list = []

    dy13_list = []
    dy32_list = []

    ep13_list = []
    ep32_list = []
    for x2 in x2_range:

        # 计算出另外两个变位系数
        x3 = x2 - x_segma_32  # x2 - x3 = x_segma_32
        x1 = x_segma_13 - x3  # x1 + x3 = x_segma_13

        with HiddenPrints(True):  # 是否打印类对象的信息，由于遍历太多，一般不打印
            # 初始化一个行星齿轮
            pt = PlanetaryGear(m, z1, z3, z2, x1, x3, x2,
                               basegear, hidelog=True)

            # 计算行星齿轮的各类参数，并验证各项约束条件是否满足，返回ret为True/False
            ret = pt.start(0.25, 0.25, xiubu=xiubuflag)

        if ret or True:
            x2_list.append(x2)
            x3_list.append(x3)
            x1_list.append(x1)

            # 外啮合副的滑动率

            yeta13_1 = pt.gear_ac.yeta1
            yeta13_3 = pt.gear_ac.yeta2
            dy13 = abs(yeta13_1-yeta13_3)
            yeta13_1_list.append(yeta13_1)
            yeta13_3_list.append(yeta13_3)
            dy13_list.append(dy13)

            # 内啮合副的滑动率
            yeta32_3 = pt.gear_cb.yeta1
            yeta32_2 = pt.gear_cb.yeta2
            dy32 = abs(yeta32_2-yeta32_3)
            yeta32_2_list.append(yeta32_2)
            yeta32_3_list.append(yeta32_3)
            dy32_list.append(dy32)

            # 内外啮合副的重合度
            ep13 = pt.gear_ac.ContactRatio
            ep32 = pt.gear_cb.ContactRatio
            ep13_list.append(ep13)
            ep32_list.append(ep32)

    # 画图
    plot2d_flag = 0

    if plot2d_flag:  # 滑动率2d
        row = 1
        col = 2
        plt.figure(figsize=(9, 5))

        plt.subplot(row, col, 1)
        plt.plot(x2_list, yeta13_1_list, label="SlidingRatio_ac_sun")
        plt.plot(x2_list, yeta13_3_list, label="SlidingRatio_ac_planet")
        plt.plot(x2_list, dy13_list, label="Delta_SlidingRatio_ac")
        plt.xlabel("x_internal")
        plt.ylabel("sliding ratio")
        plt.title("SRatio Sun&Planet,alpha_ac=%.3f,alpha_bc=%.3f" %
                  (math.degrees(alpha_work_13), math.degrees(alpha_work_32)))
        plt.legend()
        plt.grid()

        plt.subplot(row, col, 2)
        plt.plot(x2_list, yeta32_2_list, label="SlidingRatio_cb_internal")
        plt.plot(x2_list, yeta32_3_list, label="SlidingRatio_cb_planet")
        plt.plot(x2_list, dy32_list, label="Delta_SlidingRatio_cb")
        plt.xlabel("x_internal")
        plt.ylabel("sliding ratio")
        plt.title("SRatio Internal&Planet")
        plt.legend()
        plt.grid()
        plt.savefig("pic/SlidingRatio_x2.png", dpi=300, bbox_inches='tight')
        plt.show()

    if plot2d_flag:  # 重合度2d
        row = 1
        col = 1
        plt.figure(figsize=(5, 5))

        plt.subplot(row, col, 1)
        plt.plot(x2_list, ep13_list, label="ContactRatio_ac")
        plt.plot(x2_list, ep32_list, label="ContactRatio_cb")
        plt.xlabel("x_internal")
        plt.ylabel("contact ratio")
        plt.title("Contact Sun&Planet")
        plt.legend()
        plt.grid()
        plt.savefig("pic/ContactRatio_x2.png", dpi=300, bbox_inches='tight')
        plt.show()

    # pic2 : 3d
    # 给定两个变位系数
    hh2 = 0.01
    x2_range = np.arange(0.7, 1, hh2)
    x3_range = np.arange(0.4, 0.85, hh2)

    x2_list = x2_range
    x3_list = x3_range
    X2, X3 = np.meshgrid(x2_list, x3_list)

    Dy1 = np.zeros(np.shape(X2))
    Dy2 = np.zeros(np.shape(X2))
    X2ret = []
    X3ret = []
    Dy1ret = []
    Dy2ret = []
    X2zero1 = []
    X3zero1 = []
    X2zero2 = []
    X3zero2 = []
    Dy1zero = []
    Dy2zero = []
    for i in range(len(x2_range)):  # 列数不变
        x2 = x2_range[i]
        for j in range(len(x3_range)):  # 行数不变
            x3 = x3_range[j]
            # 计算内啮合变位系数和
            x_segma_32 = x2 - x3

            # 无侧隙啮合方程，计算出内啮合副的啮合角
            alpha_work_32 = inv_1(
                x_segma_32*2*math.tan(alpha)/z_segma_32+inv(alpha))

            # 中心距一致方程，计算外啮合啮合角
            alpha_work_13 = math.acos(
                z_segma_13/z_segma_32*math.cos(alpha_work_32))

            # 无侧隙啮合方程，计算外啮合副的变位系数和
            x_segma_13 = z_segma_13 * \
                (inv(alpha_work_13)-inv(alpha))/(2*math.tan(alpha))

            # 计算出太阳轮变位系数
            x1 = x_segma_13 - x3

            with HiddenPrints(True):
                # 初始化行星齿轮
                pt2 = PlanetaryGear(m, z1, z3, z2, x1, x3,
                                    x2, basegear, hidelog=True)

                # 计算校核行星齿轮
                ret = pt2.start(0.25, 0.25, xiubu=xiubuflag)
            dy1 = abs(pt2.gear_ac.yeta1-pt2.gear_ac.yeta2)
            dy2 = abs(pt2.gear_cb.yeta1-pt2.gear_cb.yeta2)
            ep1 = pt2.gear_ac.ContactRatio
            ep2 = pt2.gear_cb.ContactRatio
            Dy1[j, i] = dy1

            Dy2[j, i] = dy2
            if ret:
                pass
            else:
                X2ret.append(x2)
                X3ret.append(x3)
                Dy1ret.append(dy1)
                Dy2ret.append(dy2)

            if dy1 <= 0.01:
                X2zero1.append(x2)
                X3zero1.append(x3)
                Dy1zero.append(dy1)

            if dy2 <= 0.001:
                X2zero2.append(x2)
                X3zero2.append(x3)
                Dy2zero.append(dy2)

    plot3dflag = 1

    if plot3dflag:  # 外啮合的滑动率差3d
        fig = plt.figure(figsize=(14, 8))
        ax = fig.gca(projection='3d')
        surf = ax.plot_surface(X2, X3, Dy1, rstride=1,cstride=1, cmap='rainbow', linewidth=0, shade=False)

        #ax.scatter3D(X2zero1, X3zero1, Dy1zero, c='b', marker='s')
        ax.plot(X2zero1, X3zero1, Dy1zero, 'yellow')
        ax.scatter3D(X2ret, X3ret, Dy1ret, c='r', marker='o')
        ax.view_init(40, -14)
        ax.set_xlabel("X2")
        ax.set_ylabel("X3")
        ax.set_title("Dy1")
        #fig.colorbar(surf, shrink=0.5, aspect=5)
        #plt.savefig("pic/dy1_3d.png", dpi=500, bbox_inches='tight')
        plt.show()

    if plot3dflag:  # 内啮合的滑动率差3d
        fig = plt.figure(figsize=(14, 8))
        ax = fig.gca(projection='3d')
        surf = ax.plot_surface(X2, X3, Dy2, rstride=1,
                               cstride=1, cmap='viridis', linewidth=0, shade=False)

        ax.plot3D(X2zero2, X3zero2, Dy2zero, 'yellow')
        ax.scatter3D(X2ret, X3ret, Dy2ret, c='r', marker='o')
        ax.scatter3D(X2zero2, X3zero2, Dy2zero, c='b', marker='s')
        ax.view_init(40, -12)
        ax.set_xlabel("X2")
        ax.set_ylabel("X3")
        ax.set_title("Dy2")
        #fig.colorbar(surf, shrink=0.5, aspect=5)
        #plt.savefig("pic/dy2_3d.png", dpi=500, bbox_inches='tight')
        plt.show()

    if plot3dflag:  # 滑动率差的零线对比 2d
        fig = plt.figure(figsize=(5, 5))

        plt.plot(X2zero2, X3zero2, 'red',
                 label='Delta SlidingRatio=zero<PLANET & INTERNAL>')
        plt.plot(X2zero1, X3zero1, 'green',
                 label='Delta SlidingRatio=zero<PLANET & SUN>')
        plt.xlabel("x2")
        plt.ylabel("x3")
        plt.title("Zero Line of Delta Sliding Ratio")
        plt.legend()
        plt.grid()
        #plt.savefig("pic/Zeroline_2d.png", dpi=200, bbox_inches='tight')
        plt.show()
